from pydantic import BaseModel, Field
from langchain.agents import create_agent
from langchain.tools import tool, ToolRuntime

from model.Ark import ArkModel

def tool_access_runtime():
    class CustomerContextSchema(BaseModel):
        user_id: str = Field(description='用户的userid')

    @tool
    def get_weather(add: str, runtime: ToolRuntime) -> str:
        """ 查询一个地区的天气情况
            :Args
                add: 地理位置的名称
        """
        print(add, runtime.context.user_id)
        return '晴天'

    agent = create_agent(
        model=ArkModel().model,
        tools=[get_weather],
        context_schema=CustomerContextSchema
    )

    response = agent.invoke({
        "messages": [
            { 'role': 'user', 'content': '今天深圳的天气怎么样?' }
        ]
    }, context={"user_id": '123123123'})

    print(response['messages'][-1].content)


def run():
    tool_access_runtime()